import os
import pickle
from math import floor

from loguru import logger

from problems.problem import Problem
from websocket.msghandler import MsgContext, MessageType


class Algorithm:
    simple_callback = None

    def __init__(self,
                 pop_size: int,
                 problem: Problem,
                 kwargs: dict
                 ):
        self.parse_params(kwargs)
        self.pop_size = pop_size
        problem.__setattr__("pop_size", pop_size)
        self.problem = problem
        self.multiple_pops = []
        self.single_pops = []

    def parse_params(self, kwargs):
        """
        将参数设置到self上
        :param kwargs:
        :return:
        """
        for k, v in kwargs.items():
            self.__setattr__(k, v)

    def final_output(self):
        """
        用于决定每次迭代保存的种群体,默认是self.population
        :return:
        """
        return self.population

    def terminated(self, save_size, is_save_data):
        population = self.final_output()
        dec = population.dec
        obj = population.obj
        pop = {'dec': dec, 'obj': obj, 'FE': self.problem.fun_eval}
        if is_save_data:
            try:
                index = floor((save_size - 1) * self.problem.fun_eval / self.problem.max_fun_eval)
                self.multiple_pops[index] = pop
            except IndexError:
                self.multiple_pops.append(pop)
        else:
            self.simple_callback(pop, self.problem, self.problem.fun_eval)
            self.single_pops.append(pop)
        return self.is_terminated()

    def is_terminated(self):
        return self.problem.fun_eval >= self.problem.max_fun_eval

    # 算法的主要迭代步骤
    def each_iteration(self):
        ...

    # 算法开始运行
    def run(self, nth_run=0, problem_index=0, algorithm_index=0, save_size=0, folder_path=None, file_path=None,
            is_save_data=True):
        """
        算法的执行入口
        :param nth_run: 表示第几次执行该算法
        :param algorithm_index: 表格中的列索引
        :param problem_index:  表格中的行索引
        :param save_size: 保存算法中间种群的次数
        :param folder_path: 文件夹名字
        :param file_path:算法实验数据保存路径
        :param is_save_data: 是否保存算法结果
        :return:
        """
        if not is_save_data:
            MsgContext.callback({
                'messageType': MessageType.SIMPLE_TASK_START.value,
                'PF': self.problem.optimal_solutions.tolist(),
            })
        while not self.terminated(save_size, is_save_data):
            self.each_iteration()
        # 创建该路径下的文件,makedirs的exist_ok=True参数用于确保如果文件夹已存在，不会引发异常。
        if is_save_data:
            self.save_data(file_path, folder_path)
            FE = [item['FE'] for item in self.multiple_pops]
            return {
                'intermediate_pops': self.multiple_pops,
                'problem_index': problem_index,
                'algorithm_index': algorithm_index,
                'nth_run': nth_run,
                'file_path': file_path,
                'FE': FE
            }
        else:
            MsgContext.single_populations.append(self.single_pops)

    def save_data(self, file_path, folder_path):
        os.makedirs(folder_path, exist_ok=True)
        with open(file_path, 'wb') as my_file:
            logger.info(f"create file | {file_path}")
            # 将中间种群和性能指标一起保存到二进制文件中
            pickle.dump([self.multiple_pops, {}], my_file)
